#-*- coding: UTF-8 -*-

import redis 
import numpy as np

class FaceDatabase:
    def __init__(self, redis_port, thres):
        redis_host = '127.0.0.1'
        self.r1 = redis.Redis(host=redis_host, port=redis_port, db=1)
        self.r2 = redis.Redis(host=redis_host, port=redis_port, db=2)
        self.r3 = redis.Redis(host=redis_host, port=redis_port, db=3)
        self.thres = thres
        self.save_num = 10000
        
    def add_face(self, dbname, key, feature):
        feature = np.array(feature, dtype=np.float32)
        return self.r1.hset(dbname, key, feature.tobytes())

    def has_face(self, dbname, key):
        return self.r1.hexists(dbname, key)

    def del_face(self, dbname, key):
        return self.r1.hdel(dbname, key)
    
    def add_db(self, dbname):
        return self.r3.sadd('db', dbname)
    
    def has_db(self, dbname):
        return self.r3.sismember('db', dbname)

    def del_db(self, dbname):
        for k in self.r1.hkeys(dbname):
            self.r1.hdel(dbname, k)
        return self.r3.srem('db', dbname)

    def top1(self, dbname, f, add_stranger=0):
        keys = []
        pipe = self.r1.pipeline()

        for k in self.r1.hkeys(dbname):
            keys.append(k.decode())
            pipe.hget(dbname, k)

        if len(keys):
            features = pipe.execute()
            features = np.array([np.frombuffer(f, dtype=np.float32) for f in features])

            dist = np.dot(f, features.T)
            idx = np.argmax(dist)

            name = keys[idx]
            d = float(dist[idx])

            if d > self.thres:
                return (name, d)
        else:
            return ('none', -3.0)
            
        if not add_stranger:
            return ('notfound', -2.0) 

        keys = []
        pipe = self.r2.pipeline()

        keys_arr = [int(key) for key in self.r2.keys()]
        max_num = max(keys_arr) if len(keys_arr) else 0

        del_key_arr = [] 
        for k in self.r2.keys():
            key = k.decode()
            keys.append(key)
            pipe.get(k)

            if max_num-self.save_num >= int(key):
                del_key_arr.append(key)

        if len(keys):
            features = pipe.execute()
            features = np.array([np.frombuffer(f, dtype=np.float32) for f in features])

            dist = np.dot(f, features.T)
            idx = np.argmax(dist)

            stranger_id = keys[idx]
            name = f'stranger_{stranger_id}'
            d = float(dist[idx])

            if d > self.thres:
                return (name, d)

        f = np.array(f, dtype=np.float32)
        keys_arr = [int(key) for key in self.r2.keys()]
        max_num = max(keys_arr) if len(keys_arr) else 0

        stranger_id = max_num + 1
        name = f'stranger_{stranger_id}'
        self.r2.set(stranger_id, f.tobytes())

        for del_key in del_key_arr:
            self.r2.delete(del_key) 

        return (name, 1.0) 
